import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from torch.optim import AdamW  # 修改导入语句
from sklearn.model_selection import train_test_split

# 定义数据集类
class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

# 示例数据
texts = [
    "这是一条积极的评论，非常棒！",
    "这东西太差劲了，一点都不好用。",
    "这个电影很精彩，值得一看。",
    "服务态度太差，以后不会再来了。"
]
labels = [1, 0, 1, 0]  # 1表示积极，0表示消极

# 划分训练集和测试集
train_texts, test_texts, train_labels, test_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)

# 加载RoBERTa分词器和模型
model_name = r'D:\VS_workspace\A002_AIPython学习\Models_cache\transformers_roberta_base_chinese_cluecorpussmall'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 创建数据集和数据加载器
train_dataset = TextClassificationDataset(train_texts, train_labels, tokenizer, max_length=128)
test_dataset = TextClassificationDataset(test_texts, test_labels, tokenizer, max_length=128)

train_dataloader = DataLoader(train_dataset, batch_size=2, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=2, shuffle=False)

# 定义优化器和设备
optimizer = AdamW(model.parameters(), lr=2e-5)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 训练模型
num_epochs = 3
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch in train_dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        total_loss += loss.item()

        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch + 1}, Loss: {total_loss / len(train_dataloader)}')

# 评估模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for batch in test_dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        outputs = model(input_ids, attention_mask=attention_mask)
        logits = outputs.logits
        predictions = torch.argmax(logits, dim=1)

        total += labels.size(0)
        correct += (predictions == labels).sum().item()

print(f'Accuracy: {correct / total}')


"""
数据集类
TextClassificationDataset类继承自torch.utils.data.Dataset ， 用于封装文本数据和标签，并将文本转换为模型所需的输入格式。
数据划分
使用sklearn.model_selection.train_test_split函数将数据集划分为训练集和测试集。
模型加载
使用AutoTokenizer.from_pretrained和AutoModelForSequenceClassification.from_pretrained分别加载 RoBERTa 分词器和用于文本分类的模型。
数据加载器
使用torch.utils.data.DataLoader创建训练集和测试集的数据加载器 ， 方便批量处理数据。
训练过程
在训练过程中，将模型设置为训练模式，遍历训练数据加载器，计算损失并进行反向传播和参数更新。
评估过程
在评估过程中，将模型设置为评估模式，遍历测试数据加载器，计算预测结果并统计准确率。

注意事项
模型选择：  根据具体任务和数据情况选择合适的 RoBERTa 模型， 如中文任务可选择uer/roberta-base-chinese-cluecorpussmall 。
超参数调整： lr （ 学习率 ） 、batch_size （ 批量大小 ） 和num_epochs （ 训练轮数 ） 等超参数可能需要根据实际情况进行调整，以获得更好的性能。

"""